#该问题发生于我需要在yolo中加入差分图，需要对原始文件xml加入新的标签
import argparse
from pathlib import Path
import cv2
import xml.etree.ElementTree as ET

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="根据键盘内容改变xml内容")
    parser.add_argument('--source_img_dir', type=str, default=r"D:\tempDownload\dataset\elevator\train_backup\img\images")
    parser.add_argument('--source_xml_dir', type=str, default=r"D:\tempDownload\dataset\elevator\train_backup\img\Annotations")
    return parser.parse_args()

def Creat_Window(image_path, resize=True, source_xml_dir=""):
    xml_path = source_xml_dir / Path(image_path).with_suffix('.xml').name
    tree = ET.parse(xml_path)
    root = tree.getroot()
    background_number = root.find('backgroundNumber')

    cv2.namedWindow(f'{image_path.name}', cv2.WINDOW_NORMAL)
    cv2.moveWindow(f'{image_path.name}', 600, 250)
    image = cv2.imread(str(image_path))
    if background_number != None:
        cv2.putText(image, f"{background_number.text}", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    else:
        cv2.putText(image, f"-2", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    if resize:
        cv2.resizeWindow(f'{image_path.name}', 1500, 980)
    else:
        width = len(image[0])
        height = len(image)
        cv2.resizeWindow(f'{image_path.name}', width, height)

    cv2.imshow(f'{image_path.name}', image)

    return image, image_path


def prettyXml(element, indent, newline, level = 0): # elemnt为传进来的Elment类，参数indent用于缩进，newline用于换行
    if element:  # 判断element是否有子元素
        if element.text == None or element.text.isspace(): # 如果element的text没有内容
            element.text = newline + indent * (level + 1)
        else:
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)
    #else:  # 此处两行如果把注释去掉，Element的text也会另起一行
        #element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level
    temp = list(element) # 将elemnt转成list
    for subelement in temp:
        if temp.index(subelement) < (len(temp) - 1): # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致
            subelement.tail = newline + indent * (level + 1)
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个
            subelement.tail = newline + indent * level
        prettyXml(subelement, indent, newline, level = level + 1) # 对子元素进行递归操作

def Handle_key_Press(key, index, total_images, image, image_path, source_xml_dir):
    xml_path = source_xml_dir / Path(image_path).with_suffix('.xml').name
    tree = ET.parse(xml_path)
    root = tree.getroot()
    background_number = root.find('backgroundNumber')
    small_number = [48, 49, 50, 51, 52, 53, 54, 55, 56, 57]

    print(key)
    if key == 27:  # 如果按下ESC，退出
        exit()
    elif key in small_number and background_number == None:
        new_item = ET.SubElement(root,'backgroundNumber')
        new_item.text = str(key - 48)
    elif key in small_number and background_number != None:
        background_number.text = str(key - 48)
    elif key == 46 and background_number == None:
        new_item = ET.SubElement(root,"backgroundNumber")
        new_item.text = str(-1)
    elif key == 46 and background_number != None:
        background_number.text = str(-1)
    if key == ord('d') or key == 2555904: #右方向键对应的ASCALL码值
        index += 1
        if index >= total_images:
            index = 0  # 回到第一张图片
            print("遍历到最后一张图片")
    elif key == ord('a') or key == 2424832:  #左方向键对应的ASCALL码值
        index -= 1
        if index < 0:
            index = total_images - 1  # 跳转到最后一张图片
            print("遍历到第一张图片")
    elif key == ord('e'):
        index += 10
        if index >= total_images:
            index = 0
            print("遍历到最后一张图片")
    elif key == ord('q'):
        index -= 10
        if index < 0:
            print("遍历到第一张图片")
            index = 0
    prettyXml(root, '    ', '\n')            #执行美化方法
    # ET.dump(root)
    tree.write(xml_path)
    return index

def main():
    args = Parse_Arguments()
    source_img_dir = Path(args.source_img_dir).glob('*.jpg')
    source_xml_dir = Path(args.source_xml_dir)
    files = [f for f in source_img_dir]
    files_length = len(files)
    index = 0
    while True:
        #读取图片并显示
        image, image_path = Creat_Window(files[index], True, source_xml_dir)
        key=cv2.waitKeyEx(0)
        new_index = Handle_key_Press(key, index, files_length, image, image_path, source_xml_dir)
        if index != new_index:
            cv2.destroyAllWindows()
        #更新下标
        index = new_index
if __name__ == "__main__":
    main()
